import numpy as np
import easydict
import datetime
import pprint
import os
import shutil

import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import DataLoader

from Yolo.darknet import Darknet19
from DataSet.TongueDS import TongueDS,DLTransform
from Tools.tools import AverageMeter

from Yolo.utils.yolo import postprocess,draw_detection
from Yolo.config import cfg

def save_img(img,name='img'):
    imgpath = '/mnt/md0/Qiu/tmp/'+name+'.jpg'
    ok = cv2.imwrite(imgpath,img)
    print('write to: {} {}'.format(imgpath,ok))

try: from cv2 import cv2
except: import cv2

test_dataset  = TongueDS('test')
test_dl  = DataLoader(test_dataset, batch_size=8,shuffle=True,drop_last=True,num_workers=32)

train_dataset  = TongueDS('train')
train_dl  = DataLoader(test_dataset, batch_size=8,shuffle=True,drop_last=True,num_workers=32)

SAVED_MODEL = '/mnt/md0/Qiu/ExWorkspace/YOLO_TEST/YoloDetection_91_checkpoint.pth.tar'

model = Darknet19().cuda()
model.load_state_dict(torch.load(SAVED_MODEL)['state_dict'])
model.eval()

print('LOAD MODEL DONE!')

for _,raw_input in enumerate(test_dl):
    break

raw_input = DLTransform(raw_input,volatile=True)
raw_output = model(*raw_input)

# bbox_pred, iou_pred, prob_pred = raw_output
# bbox_pred = bbox_pred.data.cpu().numpy()
# iou_pred = iou_pred.data.cpu().numpy()
# prob_pred = prob_pred.data.cpu().numpy()


def DRAW_K(raw_output,K=0):

    bbox_pred, iou_pred, prob_pred = raw_output

    bbox_pred = bbox_pred.data.cpu().numpy()
    iou_pred = iou_pred.data.cpu().numpy()
    prob_pred = prob_pred.data.cpu().numpy()

    bbox_pred = bbox_pred[K][np.newaxis,:]
    iou_pred = iou_pred[K][np.newaxis,:]
    prob_pred = prob_pred[K][np.newaxis,:]

    bboxes, scores, cls_inds = postprocess(bbox_pred,iou_pred,prob_pred,(320,320),cfg,0.1)

    image = raw_input[0][K].cpu().data.numpy()
    image = np.transpose(image,(1,2,0))
    image = image.astype(np.uint8)

    cv2.imwrite('/mnt/md0/Qiu/tmp/t1.jpg',image)
    image = cv2.imread('/mnt/md0/Qiu/tmp/t1.jpg')

    # im2show = draw_detection(image,bboxes,scores,cls_inds,cfg,0.1)

    for box in bboxes:
        image = cv2.rectangle(image,(box[0],box[1]),(box[2],box[3]),(0,255,0),5)

    cv2.imwrite('/mnt/md0/Qiu/tmp/t1.jpg',image)

DRAW_K(raw_output,6)

# bboxes
#
# box = bboxes[0]
# box
# image = cv2.rectangle(image,(box[0],box[1]),(box[2],box[3]),(0,255,0),5)

# cv2.imwrite('/mnt/md0/Qiu/tmp/t1.jpg',image)


# with open('/mnt/md0/Qiu/tmp/R2.pkl','wb') as f:
#     D = dict()
#     D['raw_input'] = raw_input
#     D['raw_output'] = raw_output
#     import pickle
#     pickle.dump(D,f)


with open('/mnt/md0/Qiu/tmp/R1.pkl','rb') as f:
    import pickle
    D = pickle.load(f)

raw_output = D['raw_output']
raw_input = D['raw_input']
